package com.sunshine.decision;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import lombok.Data;

/**
 * Created by hadoop on 17/10/31.
 * 决策树算法
 */
public class Trees {

    public static Object[][] createDataSet() {

        Object[][] dataSet = {{1, 1, 1, 0}, {1, 1, 0, 0}, {1, 1, -1, 0}, {1, 0, 1, 0}, {0, 0, 1, 0}, {0, 1, -1, 0},
            {0, 1, 0, 0}, {0, 1, 1, 1}};
        return dataSet;
    }

    /**
     * @param dataSet 采用二维数组表示，最后一位是label
     * @return
     */
    public static double calcEntropy(Object[][] dataSet) {
        double result = 0;
        int columnSize = dataSet[0].length;
        int rows = dataSet.length;
        Map<Object, AtomicInteger> maps = Maps.newHashMap();
        for (int i = 0; i < rows; i++) {
            Object label = dataSet[i][columnSize - 1];
            if (!maps.containsKey(label)) {
                maps.put(label, new AtomicInteger(0));
            }
            maps.get(label).incrementAndGet();
        }

        for (Map.Entry<Object, AtomicInteger> entry : maps.entrySet()) {
            double prop = 1.0 * entry.getValue().get() / rows;
            result += -1 * prop * (Math.log(prop) / Math.log(2));

        }
        return result;
    }

    /**
     * 对数据集划分，按照指定的特性列的特征值做划分依据
     *
     * @param dataSet
     * @param axis
     * @param value
     * @return
     */
    public static Object[][] splitDataSet(Object[][] dataSet, int axis, Object value) {
        int columnSize = dataSet[0].length;
        int rows = dataSet.length;

        List<Object[]> lists = Lists.newArrayList();
        for (int i = 0; i < rows; i++) {

            if (!dataSet[i][axis].equals(value)) {
                continue;
            }
            Object[] row = new Object[columnSize-1];
            int index = 0;
            for (int j = 0; j < columnSize; j++) {
                if(j==axis){
                    continue;
                }
                row[index] = dataSet[i][index];
                index++;
            }
            lists.add(row);
        }

        Object[][] retDataSet = new Object[lists.size()][];
        int i = 0;
        for (Object[] row : lists) {
            retDataSet[i++] = row;
        }

        return retDataSet;
    }

    public static int chooseBestFeature(Object[][] dataSet) {
        int bestFeature = -1;
        double baseEntropy = calcEntropy(dataSet);
        int rows = dataSet.length;
        int featureSize = dataSet[0].length - 1;
        double infoGain = -1;
        double bestInfoGain = -1;
        for (int i = 0; i < featureSize; i++) {
            //处理每一列feature
            double newEntropy = 0.0;
            Map<Object, AtomicInteger> maps = Maps.newHashMap();
            for (int j = 0; j < rows; j++) {//计算当前列的唯一数值
                if (!maps.containsKey(dataSet[j][i])) {
                    maps.put(dataSet[j][i], new AtomicInteger(0));
                }
                maps.get(dataSet[j][i]).incrementAndGet();
            }
            for (Object featVa : maps.keySet()) {
                Object[][] subDataSet = splitDataSet(dataSet, i, featVa);
                double prop = maps.get(featVa).get() * 1.0 / rows;
                newEntropy += prop * calcEntropy(subDataSet);
            }
            infoGain = baseEntropy - newEntropy;
            if (infoGain > bestInfoGain) {
                bestFeature = i;
                bestInfoGain = infoGain;
            }
        }

        return bestFeature;
    }

    public static  String vote(Object[][] dataSet) {
        Map<Object,AtomicInteger> maps = Maps.newHashMap();

        int rows = dataSet.length;

        for(int i=0;i<rows;i++){
            if(!maps.containsKey(dataSet[i][0])){
                maps.put(dataSet[i][0],new AtomicInteger(0));
            }
            maps.get(dataSet[i][0]).incrementAndGet();
        }

        Object maxValue  =null;
        int max = -1;

        for(Map.Entry<Object,AtomicInteger> entry:maps.entrySet()){
            if(entry.getValue().get()>max){
                max = entry.getValue().get();
                maxValue = entry.getKey();
            }
        }
        return maxValue.toString();
    }

    public static Object createDecisionTree(Object[][] dataSet, List<String> labels) {

        int rows = dataSet.length;
        int labelIndex = dataSet[0].length - 1;

        Set<Object> sets = Sets.newHashSet();
        for (int i = 0; i < rows; i++) {
            sets.add(dataSet[i][labelIndex]);
        }

        if (sets.size() == rows) {
            return dataSet[0][labelIndex];
        }

        if (dataSet[0].length == 1) {
            return vote(dataSet);
        }

        int bestFeature = chooseBestFeature(dataSet);

        List<String> copyLabels = Lists.newArrayList(labels);

        String bestFeatureLabel = copyLabels.get(bestFeature);
        copyLabels.remove(bestFeature);
        Set<Object> bestFeatureValues = Sets.newHashSet();

        for (int i = 0; i < rows; i++) {
            bestFeatureValues.add(dataSet[i][bestFeature]);
        }

        DecisionTree decisionTree = new DecisionTree();
        decisionTree.setAttributeName(bestFeatureLabel);

        for (Object value : bestFeatureValues) {
            List<String> subFeatureLabels = Lists.newArrayList(copyLabels);
            decisionTree.children.put(value.toString(),
                createDecisionTree(splitDataSet(dataSet, bestFeature, value), subFeatureLabels));
        }

        return decisionTree;

    }

    public static void main(String[] args) {
        System.out.println(calcEntropy(createDataSet()));
        Object result = createDecisionTree(createDataSet(),Lists.newArrayList("年龄是否大于30","身高是否大于170","收入情况","yes/no"));
        System.out.println(result);
    }


    @Data
    static class DecisionTree{
        private String attributeName;
        public Map<String, Object> children = Maps.newHashMap();
    }

}
